import cv2 as cv
import npfnn
from zhnnet_npfnn import ZhnNet1
import time
import numpy as np

print('npfnn test.')
model1 = ZhnNet1()
# npfnn.save(model1.state_dict(), 'model1')
model1.load_state_dict(npfnn.load('model1'))
image_origin = cv.imread('npfnntestimg.jpg')
assert image_origin is not None, 'Image does not exit.'
image = image_origin.transpose(2, 0, 1) / 128 - 1
t1 = time.time()
predict = model1(image)
print(f'time consumption: {time.time() - t1}')
print(f'answer: {sum(predict.flatten())}')
